import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import DBSCAN
import matplotlib.patches as patches
import pickle


# 设置参数
center = (0, 0)  # 圆形区域中心坐标
radius = 20  # 圆形区域半径
num_clusters = 20  # 母点（簇中心）数量
lambda_p = 0.5  # 母点过程的强度
mean_cluster_size = 8  # 平均簇大小
gamma_shape = 2  # 用于生成簇半径的伽马分布的形状参数


# 从文件中读取之前保存的Matern簇数据，并同时初始化点是否带正方形的标记
with open("matern_cluster_data.pkl", "rb") as f:
    points_data = pickle.load(f)
    points = []
    # print(points_data[0][0][0])
    for i in range(len(points_data[0]) - 1):
        point = {
            "x": points_data[0][i][0],
            "y": points_data[0][i][1],
            "has_square": np.random.rand() < 0.2,  # 以0.2概率决定是否带正方形
        }
        points.append(point)
    parent_points_x, parent_points_y = points_data[1], points_data[2]

distances_sorted = []
for point in points:
    if point["has_square"]:
        # 计算欧几里得距离
        distance = np.sqrt(
            (point["x"] - center[0]) ** 2 + (point["y"] - center[1]) ** 2
        )
        distances_sorted.append(distance)

# print(distances_sorted)
distances_sorted = np.sort(distances_sorted)
# 创建一个包含两个子图的窗口
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))


# 找出圆内的点
in_circle = [
    (point["x"] - center[0]) ** 2 + (point["y"] - center[1]) ** 2
    <= distances_sorted[1] ** 2
    for point in points
]
print(distances_sorted[1])
circle = plt.Circle(
    center, distances_sorted[1], fill=False, color="blue", linestyle="dashed"
)
ax1.add_patch(circle)
# 在第一个子图中绘制模拟Matern簇过程的图形
ax1.scatter(
    [point["x"] for point in points],
    [point["y"] for point in points],
    label="All Points",
)
parent_points = np.array(
    [(parent_points_x[i], parent_points_y[i]) for i in range(num_clusters)]
)
ax1.scatter(
    parent_points[:, 0],
    parent_points[:, 1],
    marker="^",
    s=50,
    c="green",
    label="Cluster Centers",
)
ax1.set_xlim(center[0] - radius, center[0] + radius)
ax1.set_ylim(center[1] - radius, center[1] + radius)
ax1.set_aspect("equal", adjustable="box")
ax1.set_title("Matern Cluster Process")
ax1.set_xlabel("X-axis")
ax1.set_ylabel("Y-axis")
ax1.legend()


# 剔除圆内的点，获取用于聚类的点
points_for_clustering = [
    points[i] for i, incircle in enumerate(in_circle) if not incircle
]


# 将 points_for_clustering 中的数据转换为适合 DBSCAN 的 numpy 数组
points_for_clustering_array = np.array(
    [[point["x"], point["y"]] for point in points_for_clustering]
)


# 使用DBSCAN进行聚类
dbscan = DBSCAN(eps=2.5, min_samples=3)  # 这里的参数可根据实际情况调整
clusters = dbscan.fit_predict(points_for_clustering_array)


# 提取不同簇的点和噪声点，注意要根据points_for_clustering的索引对应回原points的索引来处理
unique_clusters = np.unique(clusters)


cluster_sets_for_draw = {cluster: [] for cluster in unique_clusters}


cluster_sets = {
    cluster: {"square_indices": [], "no_square_indices": []}
    for cluster in unique_clusters
}
# 存储用于绘图的点
for i, cluster_label in enumerate(clusters):
    index_in_original = i
    if cluster_label != -1:
        # 根据在points_for_clustering中的索引，找到在原points中的对应点
        if points_for_clustering[i]["has_square"]:
            cluster_sets[cluster_label]["square_indices"].append(index_in_original)
        else:
            cluster_sets[cluster_label]["no_square_indices"].append(index_in_original)
    else:
        # 噪声点同样要找到在原points中的对应索引
        if points_for_clustering[i]["has_square"]:
            cluster_sets[-1]["square_indices"].append(index_in_original)
        else:
            cluster_sets[-1]["no_square_indices"].append(index_in_original)


for cluster_label, cluster_points in cluster_sets_for_draw.items():
    if cluster_label != -1:
        ax2.scatter(
            [points_for_clustering[p]["x"] for p in cluster_points],
            [points_for_clustering[p]["y"] for p in cluster_points],
            label=f"Cluster {cluster_label}",
        )
    else:
        ax2.scatter(
            [points_for_clustering[p]["x"] for p in cluster_points],
            [points_for_clustering[p]["y"] for p in cluster_points],
            label="Noise Points",
            marker="x",
            c="gray",
        )


parent_points = np.array(
    [(parent_points_x[i], parent_points_y[i]) for i in range(num_clusters)]
)
"""
ax2.scatter(
    parent_points[:, 0],
    parent_points[:, 1],
    marker="^",
    s=50,
    c="green",
    label="Cluster Centers",
)
"""
ax2.set_xlim(center[0] - radius, center[0] + radius)
ax2.set_ylim(center[1] - radius, center[1] + radius)
ax2.set_aspect("equal", adjustable="box")
ax2.set_title("Matern Cluster Process with DBSCAN Clustering")
ax2.set_xlabel("X-axis")
ax2.set_ylabel("Y-axis")
# ax2.legend(loc="upper right", fontsize="small")


# 标注willie
ax1.scatter(0, 0, marker="+", s=100, c="r", label="Willie")
ax2.scatter(0, 0, marker="+", s=100, c="r", label="Willie")


# 将圆内的点变红
# ax1.scatter(
#     [point["x"] for point in points if in_circle[points.index(point)]],
#     [point["y"] for point in points if in_circle[points.index(point)]],
#     c="red",
# )
ax2.scatter(
    [point["x"] for point in points if in_circle[points.index(point)]],
    [point["y"] for point in points if in_circle[points.index(point)]],
    c="red",
)


# 为每个簇绘制圆（新添加的逻辑）
for cluster_label, cluster_points in cluster_sets.items():
    if cluster_label != -1:
        # 处理带有正方形的点
        square_indices = cluster_points["square_indices"]

        if len(square_indices) > 0:
            square_points = np.array([points_for_clustering[i] for i in square_indices])
            x_values = [d.get("x") for d in square_points if "x" in d]
            center_x = np.mean(np.array(x_values))
            y_values = [d.get("y") for d in square_points if "y" in d]

            center_y = np.mean(np.array(y_values))
            center_tuple = (center_x, center_y)
            distances = np.sqrt(
                (x_values[:] - center_x) ** 2 + (y_values - center_y) ** 2
            )
            radius = np.max(distances)
            circle1 = plt.Circle(center_tuple, radius, fill=False, color="blue")
            # ax1.add_patch(circle1)
            circle2 = plt.Circle(center_tuple, radius, fill=False, color="blue")
            ax2.add_patch(circle2)
        # 处理不带正方形的点
        no_square_indices = cluster_points["no_square_indices"]
        if len(no_square_indices) > 0:
            no_square_points = np.array(
                [points_for_clustering[i] for i in no_square_indices]
            )
            no_square_points = np.reshape(no_square_indices, (1, -1))
            center_x = np.mean(no_square_points[:, 0])
            center_y = np.mean(no_square_points[:, 1])
            center_tuple = (center_x, center_y)
            distances = np.sqrt(
                (no_square_points[:, 0] - center_x) ** 2
                + (no_square_points[:, 1] - center_y) ** 2
            )
            radius = np.max(distances)
            circle1 = plt.Circle(center_tuple, radius, fill=False, color="orange")
            # ax1.add_patch(circle1)
            circle2 = plt.Circle(center_tuple, radius, fill=False, color="orange")
            ax2.add_patch(circle2)


# 为 square_indices 对应的点画正方形
for cluster_label, cluster_points in cluster_sets.items():
    if cluster_label != -1:
        square_indices = cluster_points["square_indices"]
        if len(square_indices) > 0:
            square_points = np.array([points_for_clustering[i] for i in square_indices])
            for point in square_points:
                rect = patches.Rectangle(
                    (point["x"] - 0.5, point["y"] - 0.5),
                    1,
                    1,
                    linewidth=1,
                    edgecolor="blue",
                    facecolor="none",
                )
                ax1.add_patch(rect)
                rect = patches.Rectangle(
                    (point["x"] - 0.5, point["y"] - 0.5),
                    1,
                    1,
                    linewidth=1,
                    edgecolor="blue",
                    facecolor="none",
                )
                ax2.add_patch(rect)


plt.show()
